Skip to content

feat(pt/dpa4): Support property fitting in DPA4/SeZM models#5587

Merged
OutisLi merged 3 commits into
deepmodeling:masterfrom
OutisLi:pr/dpa4property
Jun 27, 2026
Merged

feat(pt/dpa4): Support property fitting in DPA4/SeZM models#5587
OutisLi merged 3 commits into
deepmodeling:masterfrom
OutisLi:pr/dpa4property

Conversation

@OutisLi

@OutisLi OutisLi commented Jun 25, 2026

Copy link
Copy Markdown
Collaborator

Summary by CodeRabbit

  • New Features

    • Added support for invariant property fitting with the DPA4/SeZM model family.
    • Property outputs can now be trained and read back with dedicated labels, including intensive or non-intensive aggregation.
    • Expanded model selection to recognize the DPA4 path and property-fitting configuration.
  • Bug Fixes

    • Improved output-stat handling so non-energy outputs are normalized more consistently.
    • Added validation to prevent unsupported configuration combinations.
  • Documentation

    • Updated training guides and examples for property-based DPA4/SeZM workflows.

@coderabbitai

coderabbitai Bot commented Jun 25, 2026

Copy link
Copy Markdown
Contributor

Review Change Stack

No actionable comments were generated in the recent review. 🎉

ℹ️ Recent review info
⚙️ Run configuration

Configuration used: Repository UI

Review profile: CHILL

Plan: Pro

Run ID: 8507fbf6-fd87-4736-b950-ebd5852e2a39

📥 Commits

Reviewing files that changed from the base of the PR and between abce692 and bb35245.

📒 Files selected for processing (3)
  • doc/model/dpa4.md
  • examples/property/train/README.md
  • examples/property/train/input_dpa4.json
✅ Files skipped from review due to trivial changes (2)
  • examples/property/train/README.md
  • doc/model/dpa4.md

📝 Walkthrough

Walkthrough

The PR adds invariant property fitting support for DPA4/SeZM models, including a new property model class, factory and schema support, adjusted output-stat handling and non-conservative compute paths, plus docs, examples, and tests.

Changes

SeZM property fitting

Layer / File(s) Summary
Output-stat and compute path
deepmd/pt/model/atomic_model/sezm_atomic_model.py, deepmd/pt/model/model/sezm_model.py
SeZMAtomicModel now reads fitting-network flags for output-stat handling, and SeZMModel.core_compute() adds conservative/non-conservative execution, output conversion, and overridable Inductor options.
Property model readout
deepmd/pt/model/model/sezm_property_model.py
SeZMPropertyModel validates bridging_method, remaps property outputs, routes forward/forward_lower through the shared common paths, and exposes property-specific output metadata and helper queries.
Factory and schema wiring
deepmd/pt/model/model/__init__.py, deepmd/utils/argcheck.py
SeZM model construction and schema validation accept property fitting, export SeZMPropertyModel, constrain SeZM spin fitting types, broaden DPA4 dispatch, and update argument and loss docs.
Docs and example config
doc/model/dpa4.md, examples/property/train/README.md, examples/property/train/input_dpa4.json
The DPA4 docs and property example files describe property training inputs, the property fitting configuration, and the new example JSON.
Property tests
source/tests/pt/model/test_sezm_model.py
Tests cover property-model outputs, PropertyLoss integration, serialization, and compiled execution for SeZMPropertyModel.

Sequence Diagram(s)

sequenceDiagram
  participant Factory as get_sezm_model
  participant PropertyModel as SeZMPropertyModel
  participant CoreCompute as SeZMModel.core_compute
  participant OutputMap as fit_output_to_model_output
  Factory->>PropertyModel: construct model when fitting_net.type="property"
  PropertyModel->>CoreCompute: core_compute(..., conservative=false)
  CoreCompute->>OutputMap: convert reduced outputs without force/virial graph
  PropertyModel->>PropertyModel: translate atom_{var_name} and var_name
Loading

Estimated code review effort

🎯 4 (Complex) | ⏱️ ~60 minutes

Possibly related PRs

  • deepmodeling/deepmd-kit#5580: Touches the same SeZMModel path and extends core_compute/forward_lower interface plumbing in the same model family.

Suggested labels

Core

Suggested reviewers

  • iProzd
  • wanghan-iapcm
🚥 Pre-merge checks | ✅ 5
✅ Passed checks (5 passed)
Check name Status Explanation
Description Check ✅ Passed Check skipped - CodeRabbit’s high-level summary is enabled.
Title check ✅ Passed The title clearly summarizes the main change: adding property fitting support for DPA4/SeZM models.
Docstring Coverage ✅ Passed No functions found in the changed files to evaluate docstring coverage. Skipping docstring coverage check.
Linked Issues check ✅ Passed Check skipped because no linked issues were found for this pull request.
Out of Scope Changes check ✅ Passed Check skipped because no linked issues were found for this pull request.
✨ Finishing Touches
🧪 Generate unit tests (beta)
  • Create PR with unit tests

Thanks for using CodeRabbit! It's free for OSS, and your support helps us grow. If you like it, consider giving us a shout-out.

❤️ Share

Comment @coderabbitai help to get the list of available commands.

@coderabbitai coderabbitai Bot left a comment

Copy link
Copy Markdown
Contributor

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

Actionable comments posted: 1

🧹 Nitpick comments (1)
deepmd/pt/model/model/sezm_property_model.py (1)

121-153: 📐 Maintainability & Code Quality | 🔵 Trivial | ⚡ Quick win

Keep core_compute() signature aligned with SeZMModel.

Line 121 overrides the widened SeZMModel.core_compute() API but drops the new conservative keyword. That makes SeZMPropertyModel fail with TypeError for any generic caller that uses the base-class interface.

Proposed fix
     def core_compute(
         self,
         coord: torch.Tensor,
         atype: torch.Tensor,
         edge_index: torch.Tensor,
         edge_vec: torch.Tensor,
         edge_scatter_index: torch.Tensor,
         edge_mask: torch.Tensor,
         fparam: torch.Tensor | None = None,
         aparam: torch.Tensor | None = None,
         charge_spin: torch.Tensor | None = None,
         comm_dict: dict[str, torch.Tensor] | None = None,
         extended_atype: torch.Tensor | None = None,
         extended_coord_corr: torch.Tensor | None = None,
         embedding_only: bool = False,
+        conservative: bool = False,
     ) -> dict[str, torch.Tensor]:
         """Compute property outputs through the SeZM forward-only graph."""
         return super().core_compute(
             coord,
             atype,
             edge_index,
             edge_vec,
             edge_scatter_index,
             edge_mask,
             fparam=fparam,
             aparam=aparam,
             charge_spin=charge_spin,
             comm_dict=comm_dict,
             extended_atype=extended_atype,
             extended_coord_corr=extended_coord_corr,
             embedding_only=embedding_only,
             conservative=False,
         )
🤖 Prompt for AI Agents
Verify each finding against current code. Fix only still-valid issues, skip the
rest with a brief reason, keep changes minimal, and validate.

In `@deepmd/pt/model/model/sezm_property_model.py` around lines 121 - 153,
`SeZMPropertyModel.core_compute` is missing the new `conservative` keyword that
exists on `SeZMModel.core_compute`, so the override no longer matches the
base-class API and generic callers can hit a `TypeError`. Update `core_compute`
in `SeZMPropertyModel` to accept the same signature as `SeZMModel.core_compute`,
including `conservative`, and forward that argument through the
`super().core_compute(...)` call unchanged.
🤖 Prompt for all review comments with AI agents
Verify each finding against current code. Fix only still-valid issues, skip the
rest with a brief reason, keep changes minimal, and validate.

Inline comments:
In `@deepmd/utils/argcheck.py`:
- Around line 3263-3271: The DPA4/SeZM fitting schema is missing the supported
sezm_ener variant, causing validation to reject configs that the runtime still
accepts. Update the fitting type Variant in the schema builder around the
DPA4/SeZM arguments to include fitting_args_plugin.get_argument("sezm_ener")
alongside dpa4_ener and property, keeping the existing default_tag and doc
behavior intact so get_sezm_model() and get_sezm_spin_model() remain valid.

---

Nitpick comments:
In `@deepmd/pt/model/model/sezm_property_model.py`:
- Around line 121-153: `SeZMPropertyModel.core_compute` is missing the new
`conservative` keyword that exists on `SeZMModel.core_compute`, so the override
no longer matches the base-class API and generic callers can hit a `TypeError`.
Update `core_compute` in `SeZMPropertyModel` to accept the same signature as
`SeZMModel.core_compute`, including `conservative`, and forward that argument
through the `super().core_compute(...)` call unchanged.
🪄 Autofix (Beta)

Fix all unresolved CodeRabbit comments on this PR:

  • Push a commit to this branch (recommended)
  • Create a new PR with the fixes

ℹ️ Review info
⚙️ Run configuration

Configuration used: Repository UI

Review profile: CHILL

Plan: Pro

Run ID: 2e654586-7523-4e51-a585-094ed7f5851a

📥 Commits

Reviewing files that changed from the base of the PR and between 5733301 and 58e17c3.

📒 Files selected for processing (8)
  • deepmd/pt/model/atomic_model/sezm_atomic_model.py
  • deepmd/pt/model/model/__init__.py
  • deepmd/pt/model/model/sezm_model.py
  • deepmd/pt/model/model/sezm_property_model.py
  • deepmd/utils/argcheck.py
  • doc/model/dpa4.md
  • examples/water/dpa4/input_property.json
  • source/tests/pt/model/test_sezm_model.py

Comment thread deepmd/utils/argcheck.py
@codecov

codecov Bot commented Jun 25, 2026

Copy link
Copy Markdown

Codecov Report

❌ Patch coverage is 83.15789% with 16 lines in your changes missing coverage. Please review.
✅ Project coverage is 82.30%. Comparing base (5733301) to head (bb35245).
⚠️ Report is 5 commits behind head on master.

Files with missing lines Patch % Lines
deepmd/pt/model/model/sezm_property_model.py 79.16% 10 Missing ⚠️
deepmd/pt/model/atomic_model/sezm_atomic_model.py 84.21% 3 Missing ⚠️
deepmd/pt/model/model/__init__.py 82.35% 3 Missing ⚠️
Additional details and impacted files
@@            Coverage Diff             @@
##           master    #5587      +/-   ##
==========================================
+ Coverage   82.27%   82.30%   +0.02%     
==========================================
  Files         887      888       +1     
  Lines      100331   100533     +202     
  Branches     4060     4056       -4     
==========================================
+ Hits        82550    82744     +194     
- Misses      16320    16326       +6     
- Partials     1461     1463       +2     

☔ View full report in Codecov by Harness.
📢 Have feedback on the report? Share it here.

🚀 New features to boost your workflow:
  • ❄️ Test Analytics: Detect flaky tests, report on failures, and find test suite problems.
  • 📦 JS Bundle Analysis: Save yourself from yourself by tracking and limiting bundle sizes in JS merges.

@njzjz-bot

Copy link
Copy Markdown
Contributor

Thanks for adding the SeZM/DPA4 invariant property path. The eager forward contract, public key translation (atom_<var> / <var>), reduction handling, and serialization path look reasonable from code inspection.

One blocker before this can merge: the new compiled-property training test is currently failing in CI on both Python 3.10 and 3.13. The failing test is:

source/tests/pt/model/test_sezm_model.py::TestSeZMModelProperty::test_compile_matches_eager_and_backpropagates

The failure happens at loss.backward() after the compiled property forward, with a TorchInductor C++ compile error like:

torch._inductor.exc.InductorError: CppCompileError: C++ compile error
error: ‘decltype’ evaluates to ‘float’, which is not a class or enumeration type

Examples:

Since the PR exposes use_compile=True for SeZMPropertyModel and the test asserts compiled backprop works, this should either be fixed in the property compile/backward path, or the property model should explicitly fall back / reject use_compile=True until that path is supported. Otherwise users can hit a runtime compiler failure during training.

Reviewed by OpenClaw 2026.6.8 (844f405) (model: custom-chat-jinzhezeng-group/gpt-5.5)

@njzjz-bot

Copy link
Copy Markdown
Contributor

One more documentation/example issue I noticed: doc/model/dpa4.md points to examples/water/dpa4/input_property.json as a “complete input” for property fitting, and that input uses:

"property_name": "band_prop"

with training systems ../data/data_0, ../data/data_1, ../data/data_2 and validation ../data/data_3. However, the checked-in water example data only contains box.npy, coord.npy, energy.npy, and force.npy under those systems; there is no band_prop.npy. So the advertised complete property example does not appear runnable as-is.

Could we either add matching property-label arrays for the example data, or change the docs/example wording to make clear this is a template that requires user-supplied band_prop.npy labels?

Reviewed by OpenClaw 2026.6.8 (844f405) (model: custom-chat-jinzhezeng-group/gpt-5.5)

@wanghan-iapcm wanghan-iapcm left a comment

Copy link
Copy Markdown
Collaborator

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

please resolve njzjz's comment.

@njzjz-bot

Copy link
Copy Markdown
Contributor

Thanks for the follow-up fixes. I re-reviewed the current head (abce692) and the two earlier blockers look addressed:

  • the compiled property test failure is now fixed; the current PR checks are green, including the Python test matrix;
  • the example now points at the QM9 property subset that actually provides band_prop.npy labels, so the documented complete input is runnable.

One remaining issue I would fix before merge: the DPA4/SeZM input schema still omits sezm_ener as an allowed fitting_net.type in deepmd/utils/argcheck.py. Runtime construction accepts it in get_sezm_model() / get_sezm_spin_model() via:

if fitting_net["type"] in ("dpa4_ener", "sezm_ener"):

but the schema variant currently lists only dpa4_ener and property. That means an existing valid config using model.type: dpa4 + fitting_net.type: sezm_ener can be rejected by argument validation even though the model factory still supports it. Please add:

fitting_args_plugin.get_argument("sezm_ener"),

next to dpa4_ener in the DPA4/SeZM fitting variant.

Reviewed by OpenClaw 2026.6.8 (844f405) (model: custom-chat-jinzhezeng-group/gpt-5.5)

Comment thread deepmd/utils/argcheck.py
Comment thread examples/property/train/input_dpa4.json
@OutisLi OutisLi requested a review from njzjz June 27, 2026 02:03
@OutisLi OutisLi added this pull request to the merge queue Jun 27, 2026
Merged via the queue into deepmodeling:master with commit 1f0ca6c Jun 27, 2026
70 checks passed
@OutisLi OutisLi deleted the pr/dpa4property branch June 27, 2026 18:05
Sign up for free to join this conversation on GitHub. Already have an account? Sign in to comment

Projects

None yet

Development

Successfully merging this pull request may close these issues.

4 participants